# -*- coding: utf-8 -*-
"""
Created on Mon Apr 24 10:30:41 2023

@author: lenovo
"""

# import xarray as xr
import os
import numpy as np
import math
import scipy.interpolate
from osgeo import gdal
import re
import glob
# import shutil
from tqdm import tqdm # 运行计时

class para:
    climate_path=r'E:\drivedata\era5data'
    dem_path=r'E:\GDEM V3'
    slop_path=r'E:\slop30m'

    landtypedir = r'F:\公司项目\陕西\商洛\land30m'
    out_path=r'F:\公司项目\陕西\商洛\data\GEP'
    inputyear=2023
    lon1=-1
    lon2=-1
    lat1=-1
    lat2=-1
    imlat=0
    imlon=0
    imwidth=3601
    imheight=3601
    num='0'
################################################
def read_img(filename):
    dataset = gdal.Open(filename)  # 打开文件
    if dataset == None:
        print(filename+"文件无法打开")
        return
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    im_bands = dataset.RasterCount #波段数
    im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height).astype(np.float32)  # 将数据写成数组，对应栅格矩阵
    im_lon=[im_geotrans[0]+i*im_geotrans[1] for i in range(im_width)]
    im_lat=[im_geotrans[3]+i*im_geotrans[5] for i in range(im_height)]
    
    del dataset  # 关闭对象，文件dataset   
    return im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat
#=========================
# 保存tif文件函数
def write_img(im_data, im_geotrans, im_proj, path, nodata=None):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if (dataset != None):
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影        
    for i in range(im_bands):
        if (nodata != None):
            dataset.GetRasterBand(i + 1).SetNoDataValue(nodata)
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

#=======================   
def readasc(fname):
    res=np.loadtxt(fname,skiprows=6)
    res[res==-9999]=np.nan
    return res
#======================================
def interpolation(lat,lon,data,met):
                
    latc=para.imlat
    lonc=para.imlon
    
    lon_c,lat_c=np.meshgrid(lonc,latc)

    lon_b,lat_b=np.meshgrid(lon,lat)
    #method{‘linear’, ‘nearest’, ‘cubic’}, optional
    res = scipy.interpolate.griddata((lon_b.ravel(), lat_b.ravel()),
                                    data.ravel(), (lon_c, lat_c), 
                                    method=met,fill_value=np.nan)
    return res
#==========================
def get_ap(tp):
    if tp==101: #常绿阔叶
        s=5.75
        n=3.52
        c=11.76
    elif tp==102: 
        s=3.38
        n=2.35
        c=8.41
    elif tp==103:
        s=5.04
        n=3.52
        c=20.18
    elif tp==104:
        s=3.38
        n=2.35
        c=10.08
    elif tp==105:
        s=5.09
        n=2.46
        c=16.8
    elif tp==61:
        s=3.6
        n=2.26
        c=10.76
    elif tp==106:
        s=4.03
        n=2.64
        c=11.76
    elif tp==107:
        s=2.94
        n=1.57
        c=7.88
    elif tp==108:
        s=3.73
        n=2.35
        c=10.08
    elif tp==62:
        s=2.81
        n=1.75
        c=7.93
    elif tp==21:
        s=3.6
        n=2.56
        c=10.6
    elif tp==22:
        s=2.94
        n=1.57
        c=8.41
    elif tp==23:
        s=2.94
        n=1.57
        c=8.41
    elif tp==63:
        s=2.54
        n=1.52
        c=7.18
    elif tp==41:
        s=4.03
        n=2.75
        c=8.87
    elif tp==42:
        s=2.5
        n=1.57
        c=8.41
    elif tp==109:
        s=3.38
        n=2.56
        c=8.41
    elif tp==110:
        s=3.16
        n=2.17
        c=6.17
    elif tp==111:
        s=3.6
        n=2.26
        c=10.76
    elif tp==112:
        s=2.81
        n=1.75
        c=7.93
    elif tp==24:
        s=2.54
        n=1.52
        c=7.18
    elif tp==31:
        s=4.03
        n=1.97
        c=10.08
    elif tp==32:
        s=3.11
        n=1.52
        c=7.41
    elif tp==33:
        s=2.85
        n=1.32
        c=6.73
    elif tp>=34 and tp<=37:
        s=7.06
        n=0
        c=10.08     
    else:
        s=0
        n=0
        c=0
        
    return s,n,c
################################################    
tiffile_all = glob.glob(para.landtypedir+'\\'+'*.tif')
for fn in tqdm(tiffile_all):
    tif_file_folder,tif_file_name = os.path.split(fn)
    para.num=re.findall('\d+', tif_file_name)
    
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)

    para.imlat=im_lat
    para.imlon=im_lon
    para.imheight=im_height
    para.imwidth=im_width
    para.lat1=im_lat[0]+0.5
    para.lat2=im_lat[-1]-0.5
    para.lon1=im_lon[0]-0.5
    para.lon2=im_lon[-1]+0.5

       


    cols=im_width
    rows=im_height
    Vap=np.zeros((rows,cols))
    for i in range(0,rows):
        for j in range(0,cols):
            if im_data[i,j]==255:
                Vap[i,j]=np.nan
            else:
                s,n,c=get_ap(im_data[i,j])
                Vap[i,j]=(1263*s+1263*n+150*c)/1000/1000
    
    tif_file_name=tif_file_name.split('_')[1]
                
    Vap[Vap<0]=0
    Vap[np.isnan(Vap)]=-9999
    fout=para.out_path+'\\空气净化_'+tif_file_name
    write_img(Vap, im_geotrans, im_proj, fout, -9999)
